import FetchStream from "./FetchStream.js";

const BAiCHUAN_API = 'https://api.baichuan-ai.com/v1/chat/completions';
const API_KEY = 'sk-a63a8102be504750dc22b771041c015d';

// 最大返回token数
const MAX_TOKENS = 8192;

// 记录每次请求的返回结果,用于多轮询问
const CATCHE_DB = {
    /*
    
    'conversation_id': {
        create_time: ''
        history: [] 
    }
    */

};

export async function baichuan_stream(config, res) {

    if (!config.id) {
        res.send({
            status: 1,
            error: '缺少conversation id'
        });
        return;
    }

    if (!CATCHE_DB[config.id]) {
        CATCHE_DB[config.id] = {
            create_time: Date.now(),
            history: [
                { "role": "assistant", "content": "你现在是C/C++和Rust的资深专家，工作是将C/C++代码转换为Rust代码，在转换代码时，需要精确匹配C++的行为，要求尽可能使用Rust标准库实现，并尽可能不使用unsafe关键字，每次转换，你都需要给出完整可执行的Rust代码，不要做额外解释。" },
                // { "role": "assistant", "content": "你现在是C/C++和Rust的资深专家，工作是将C/C++代码转换为Rust代码，在转换代码时，需要精确匹配C++的行为，要求尽可能使用Rust标准库实现，并尽可能不使用unsafe关键字，直接给出代码，不要做额外解释。" },
            ]
        }
    }

    if (config.prompt) {
        CATCHE_DB[config.id].history.push({ "role": "user", "content": config.prompt })
    } else {
        CATCHE_DB[config.id].history.push({ "role": "user", "content": '请将下列代码转换为Rust代码：\n' + config.code })
    }

    const messages = CATCHE_DB[config.id].history;
    console.log(`baichuan messages[${config.id}] length: `, messages.length);

    const data = {
        model: 'Baichuan2-Turbo',
        messages,
        max_tokens: MAX_TOKENS,
        stream: true
    }

    let content = ''; // 拼接每一个分块数据
    const fetchStream = new FetchStream({
        url: BAiCHUAN_API,
        requestInit: { // fetch 配置项
            method: "POST",
            headers: {
                "Content-Type": "application/json",
                "Accept": "text/event-stream",
                "Authorization": `Bearer ${API_KEY}`
            },
            body: JSON.stringify(data),
        },
        onmessage: (dataList, _index) => {

            for (const item of dataList) {
                // console.log('------------------------');
                // console.log(item);
                const data_idx = item.indexOf('data: ');
                // const data_str = item.split("data: ")[1];
                let data_str = '';
                if (data_idx === -1) {
                    data_str = item;
                } else {
                    data_str = item.substring(data_idx + 6);
                }
                if (data_str && data_str !== '[DONE]') {
                    try {
                        const data_json = JSON.parse(data_str);
                        content += data_json.choices[0].delta.content;
                        // console.log(content);
                        try {
                            res.write(JSON.stringify({
                                status: 0,
                                content: data_json.choices[0].delta.content,
                                id: config.id,
                            }) + '\n\n');
                        } catch (e) {
                            console.error(e);
                        }
                    } catch (e) {
                        console.log(e);
                        // 将解析错误的内容返回缓存，在下一次消息中添加到新消息的头部
                        return item;
                    }
                }
            }
        },
        ondone: () => {
            // console.log('done!', content);
            // 将回答加入历史记录
            CATCHE_DB[config.id].history.push({ "role": "assistant", "content": content });
            res.write('[DONE]');
            res.end();
        },
        ontimeout: () => {
            console.log('timeout!');
        }
    });
    // res.send({
    //     status: 0,
    //     content: '```rust\n// This code is generated by Baichuan. Please do not modify it.\n```',
    // })
}

/**
 *      tool: 'c2rust',
 *      code: '',
 *      prompt: '',
 *      id: '', // 用于标识不同的的对话
*/

export async function baichuan(config) {
    console.log('[baichuan] start transcode...');

    return new Promise(async (resolve, reject) => {

        if (!config.id) return reject("id required");

        if (!CATCHE_DB[config.id]) {
            CATCHE_DB[config.id] = {
                create_time: Date.now(),
                history: [
                    { "role": "assistant", "content": "你现在是C/C++和Rust的资深专家，工作是将C/C++代码转换为Rust代码，在转换代码时，需要精确匹配C++的行为，要求尽可能使用Rust标准库实现，并尽可能不使用unsafe关键字，每次转换，你都需要给出完整可执行的Rust代码，不要做额外解释。" },
                ]
            }
        }

        if (config.prompt) {
            CATCHE_DB[config.id].history.push({ "role": "user", "content": config.prompt })
        } else {
            CATCHE_DB[config.id].history.push({ "role": "user", "content": '请将下列代码转换为Rust代码：\n' + config.code })
        }

        const messages = CATCHE_DB[config.id].history;
        console.log(`baichuan messages[${config.id}] length: `, messages.length);

        const data = {
            model: 'Baichuan2-Turbo',
            messages,
            max_tokens: MAX_TOKENS,
        }

        let answer = null;
        try {
            const response = await fetch(BAiCHUAN_API, {
                method: "POST",
                headers: {
                    "Content-Type": "application/json",
                    "Authorization": `Bearer ${API_KEY}`
                },
                body: JSON.stringify(data),
            });
    
            answer = await response.json();
        } catch(e) {
            console.error('[baichun] API 连接错误:', e);
            reject('获取数据失败。');
        }

        if (answer.error) {
            return reject(answer.error.message);
        }

        // console.log(answer.choices[0].message.content);

        // 将回答加入历史记录
        CATCHE_DB[config.id].history.push({ "role": "assistant", "content": answer.choices[0].message.content });

        const reg = /```rust\n([\s\S]*?)```/g;

        let result = '';
        const codes = [];
        while (result = reg.exec(answer.choices[0].message.content)) {
            codes.push(result[1]);
        }
        // const rust_code = answer.choices[0].message.content.match(/```rust\n([\s\S]*)```/)[1];

        // return codes.join('\n');
        resolve((codes[0] || '').trim());
    })
}